Motivation

Based on #2961
Integrate Qwen3.5 397B FP8 PTPC MOE in decoding phase for batch size 1-32 with FP8 weight decompression via hsaco.

Technical Details

Major optimization:

• Batch1: a new special kernel moe_gemm_batch1 can directly use topk_ids and topk_weight and there is no dependency on moe_sorting
• Batch 2~15: stage 1: moe_gemm_batch + stage 2: moe_2stage_splitk
• Batch16~32: stage 1: moe_gemm_batch + stage 2: moe_2stage_down_loopn
• FP8: since it's memory bound, a new kernel will do weight decompression instead of additional dynamic quantization kernels.

The function is disabled by default, should use AITER_MOE_SMALL_BATCH=1 to enable.
The corresponding source file for .co is moe.py.

Test Plan

AITER_MOE_SMALL_BATCH=1 python3 op_tests/test_moe_2stage.py -q 2 -a silu -e 512 -k 10 -dim 4096,128 -p t  -t 1

Test Result

Test results on MI308X (non-PTL)
AITER_MOE_SMALL_BATCH=1 for batch size 1/2/4/8/10/12/16/32

AITER_MOE_SMALL_BATCH=1 python3 op_tests/test_moe_2stage.py -q 2 -a silu -e 512 -k 10 -dim 4096,128 -p t  -t 1

AITER_MOE_SMALL_BATCH=0 for batch size 1/2/4/8/10/12/16/32

AITER_MOE_SMALL_BATCH=0 python3 op_tests/test_moe_2stage.py -q 2 -a silu -e 512 -k 10 -dim 4096,128 -p t  -t 1

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.